# Defines the core data structures used internally by the agent components.
# These might overlap with or extend schemas defined in src/schemas/agent.py

from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional, Literal
import uuid
from datetime import datetime, timezone

# --- Re-exporting or redefining core schemas for agent internal use ---
# Ensures agent components have clear data contracts.
from ..schemas.agent import Task, Plan, Step, AgentObservation, MemoryRecord # Use schemas defined earlier
from ..schemas.common import StatusResponse, ErrorResponse # Useful for results

# You might add agent-specific internal states here if needed,
# e.g., representing the agent's internal belief state or focus.

class AgentContext:
    """
    A helper object potentially passed around within the orchestrator loop
    to hold current state relevant to the ongoing task execution.
    (Alternatively, Orchestrator can manage state internally).
    """
    task: Task
    plan: Optional[Plan] = None
    current_observation: Optional[AgentObservation] = None
    step_history: List[Step] = [] # History of executed steps in current run
    # Add other contextual info like error count, reflection notes etc.

    def __init__(self, task: Task):
        self.task = task
        self.step_history = []

    def add_executed_step(self, step: Step):
         """Adds a step after execution, updating its status/result."""
         # Find if step with same ID exists and update, otherwise append
         found = False
         for i, s in enumerate(self.step_history):
              if s.id == step.id and s.tool_name == step.tool_name: # Check tool too?
                   self.step_history[i] = step
                   found = True
                   break
         if not found:
              self.step_history.append(step)

    def update_plan(self, plan: Plan):
         self.plan = plan
         # Potentially reset step history if plan changes significantly?

    def set_observation(self, observation: AgentObservation):
         self.current_observation = observation